import os
os.environ["HF_ENDPOINT"] = "https://hf-api.gitee.com"
from openmind_hub import set_platform
set_platform("gitee")

from openmind_hub import OM_HUB_CACHE, om_hub_download, snapshot_download
import unittest
import time
import shutil


token = "123456789()"
repo_id = "guoxiaozhen/gxzdataset1011"
author = "guoxiaozhen"
reponame = "gxzdataset1011"

def search_dir_File(pathname,reponame):
    folders = [f for f in os.listdir(pathname) if os.path.isdir(os.path.join(pathname, f))]
    filtered_folders = [f for f in folders if 'openmind' in f or 'model' in f or 'space' in f or 'dataset' in f or reponame in f]
    return filtered_folders

class TestOpenmindHub(unittest.TestCase):
    def test_01_gitee_om_hub_download_success(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="rrrrrr.json",
                subfolder=None,
                user_agent=None,
                local_dir="./",
                force_download=False,
                token=token,
                # local_files_only=True,
                repo_type="model",
                revision="master"
            )
            time.sleep(3)
            self.assertIsNotNone("./rrrrrr.json", "下载的文件不存在")
        except:
            raise Exception("下载文件有报错")
        else:
            os.remove("./rrrrrr.json")
    def test_02_gitee_om_hub_download_wrongtoken(self):
        try:
            om_hub_download(
                repo_id=repo_id,
                filename="rrrrrr.json",
                subfolder=None,
                user_agent=None,
                local_dir="./",
                force_download=False,
                token=token+"a",
                # local_files_only=True,
                repo_type="model",
                revision="master"
            )
            raise Exception("om_hub_download未校验token")
        except:
            pass
    def test_03_gitee_om_hub_download_noexistrepo(self):
        try:
            om_hub_download(
                repo_id=repo_id + "a",
                filename="rrrrrr.json",
                subfolder=None,
                user_agent=None,
                local_dir="./",
                force_download=False,
                token=token,
                # local_files_only=True,
                repo_type="model",
                revision="master"
            )
            raise Exception("om_hub_download未校验repo是否存在")
        except:
            pass
    def test_04_gitee_snapshot_download_cachedir_success(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision="master",
                cache_dir=OM_HUB_CACHE,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(OM_HUB_CACHE,"")
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            for folder in dirs:
                folder_to_delete = os.path.join(OM_HUB_CACHE, folder)
                shutil.rmtree(folder_to_delete)
    def test_05_gitee_snapshot_download_wrong_token(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision="master",
                cache_dir=OM_HUB_CACHE,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token+"a",
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件未校验TOKEN")
        except:
           pass
    def test_06_gitee_snapshot_download_noexist_repo(self):
        try:
            snapshot_download(
                repo_id=repo_id+"a",
                repo_type="dataset",
                revision="master",
                cache_dir=OM_HUB_CACHE,
                #local_dir="space_test0516",
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("下载文件未校验repo是否存在")
        except:
           pass
    def test_07_gitee_snapshot_download_localdir_success(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="dataset",
                revision="master",
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"./"+repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            time.sleep(3)
            dirs=search_dir_File(r"./"+author,reponame)
            self.assertTrue(len(dirs),"下载的文件不存在")
        except:
           raise Exception("下载文件有报错")
        else:
            shutil.rmtree(r"./"+author)
    def test_08_gitee_snapshot_download_type_id_notconsistent(self):
        try:
            snapshot_download(
                repo_id=repo_id,
                repo_type="model",
                revision="master",
                # cache_dir=OM_HUB_CACHE,
                local_dir=r"./" + repo_id,
                local_dir_use_symlinks="auto",
                library_name=None,
                library_version=None,
                user_agent=None,
                proxies=None,
                resume_download=True,
                force_download=False,
                token=token,
                local_files_only=False,
                # allow_patterns="allowed_folder\*",    # 只下载allowed_folder目录下的文件
                # ignore_patterns=["*.log", "*.txt"],    # 不下载以.log和.txt结尾的文件
                max_workers=8,
                tqdm_class=None,
                endpoint=None,
            )
            raise Exception("没有校验type与ID不一致的情况")
        except:
           pass



if __name__ == '__main__':
    unittest.main()